package com.gitee.booting.idem4j;

import com.gitee.booting.idem4j.properties.Idem4jProperties;
import com.gitee.booting.idem4j.util.JavaUtils;
import org.springframework.data.redis.core.RedisCallback;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.http.HttpStatus;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.mvc.method.RequestMappingInfo;
import org.springframework.web.servlet.mvc.method.annotation.RequestMappingHandlerMapping;

import javax.servlet.*;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import java.util.function.Function;
import java.util.regex.Pattern;

public class Idem4jFilter implements Filter {
    public static final String HEADER_TOKEN = "X-Idempotent-Token";
    public static final String MESSAGE = "Invalid idempotent token.";
    private static final Map<String, Idem4jProperties.Control> CACHE = new HashMap<>();
    private static final Function<String, String> KEY = s -> "idempotent:tokens:" + s;
    private final RequestMappingHandlerMapping mapping;
    private final Idem4jProperties properties;
    private final RedisTemplate<Object, Object> template;

    private static final Function<String, String> RESPONSE = token -> {
        if (null != token) {
            return "{\"status\":200,\"code\":\"operation.ok\",\"message\":" +
                "\"Retrieve an idempotent token success.\",\"error\":null,\"data\":\"" + token + "\"}";
        } else {
            return "{\"status\":400,\"code\":\"operation.fail\",\"message\":" +
                "\"Retrieve an idempotent token failed.\",\"error\":null,\"data\":null}";
        }
    };
    private static final long INC_STEP = 1;
    private static final int EXPIRED = 2;
    private static final byte[] INITIAL = "1".getBytes();
    private static final byte[] NOBYTES = new byte[0];

    public Idem4jFilter(RequestMappingHandlerMapping mapping, Idem4jProperties properties, RedisTemplate<Object, Object> template) {
        this.mapping = mapping;
        this.properties = properties;
        this.template = template;
    }

    @Override
    public void init(FilterConfig filterConfig) {
        // initial patterns
        Map<RequestMappingInfo, HandlerMethod> handlers = this.mapping.getHandlerMethods();
        Map<String, Pattern> patterns = new HashMap<>();
        properties.getMapping().forEach((pattern, supplier) -> {
            if (!patterns.containsKey(pattern)) {
                patterns.put(pattern, Pattern.compile(pattern));
            }
        });
        Set<Map.Entry<String, Pattern>> set = patterns.entrySet();
        for (Map.Entry<String, Pattern> e : set) {
            handlers.forEach((inf, method) -> {
                inf.getPatternsCondition().getPatterns().forEach(s -> {
                    if (e.getValue().matcher(s).matches()) {
                        Idem4jProperties.Control control = properties.getMapping().get(e.getKey());
                        if (null == control.getGenerator()) {
                            // allow empty generator to use default one
                            control.setGenerator(JavaUtils.DEFAULT.getGenerator());
                        }
                        CACHE.put(s, control);
                    }
                });
            });
        }
    }

    @Override
    public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain) throws IOException, ServletException {
        final HttpServletRequest httpRequest = (HttpServletRequest) request;
        final HttpServletResponse httpResponse = (HttpServletResponse) response;
        final String pattern = httpRequest.getRequestURI();
        if (properties.getToken().equalsIgnoreCase(pattern)) {
            String token = CACHE.getOrDefault(pattern, JavaUtils.DEFAULT).getGenerator().get();
            template.executePipelined((RedisCallback<Boolean>) connection -> {
                // for spring-boot 1.4.3 capabilities (spring-data-redis:1.7.2.RELEASE)
                byte[] key = KEY.apply(token).getBytes(StandardCharsets.UTF_8);
                connection.setNX(key, INITIAL);
                connection.expire(key, properties.getExpire());
                // must be null, otherwise will throw an exception
                return null;
            });
            httpResponse.setStatus(HttpStatus.OK.value());
            httpResponse.getWriter().write(RESPONSE.apply(token));
            return;
        } else if (CACHE.containsKey(pattern)) {
            final Idem4jProperties.Control control = CACHE.get(pattern);
            final String token = httpRequest.getHeader(HEADER_TOKEN);
            final byte[] redisKey;
            final Long count;
            if (null != token) {
                redisKey = KEY.apply(token).getBytes(StandardCharsets.UTF_8);
                count = template.execute((RedisCallback<Long>) c -> c.incrBy(redisKey, INC_STEP));
            } else {
                redisKey = NOBYTES;
                count = INC_STEP;
            }
            boolean error = false;
            if (null != count && INC_STEP == count) {
                if (this.properties.isFailOnAbsent() || control.isFailOnAbsent()) {
                    // fail on absent, remove current token (set by incr operation)
                    if (NOBYTES != redisKey) {
                        template.execute((RedisCallback<Long>) c -> c.del(redisKey));
                    }
                    error = true;
                } else {
                    if (NOBYTES != redisKey) {
                        // always expire to prevent from redis bombing
                        template.executePipelined((RedisCallback<Boolean>) c -> {
                            c.expire(redisKey, properties.getExpire());
                            // initial value is 1, so it should be 1+1 once this request is allowed
                            c.incr(redisKey);
                            return null;
                        });
                    }
                    // allow unofficial idempotent token or not
                    if (!this.properties.isAllowUnofficial()) {
                        error = true;
                    }
                }
            }
            if (null == count || count > EXPIRED) {
                // unexpected exception or repeated token usage
                error = true;
            }
            if (error) {
                httpResponse.sendError(HttpStatus.FORBIDDEN.value(), MESSAGE);
                return;
            }
        }
        chain.doFilter(request, response);
    }

    @Override
    public void destroy() {

    }
}
